{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import cv2\n",
    "import numpy as np\n",
    "import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ai format\n",
    "# 0  right_shoulder\n",
    "# 1  right_elbow\n",
    "# 2  right_wrist\n",
    "# 3  left_shoulder\n",
    "# 4  left_elbow\n",
    "# 5  left_wrist\n",
    "# 6  right_hip\n",
    "# 7  right_knee\n",
    "# 8  right_ankle\n",
    "# 9  left_hip\n",
    "# 10 left_knee\n",
    "# 11 left_ankle\n",
    "# 12 head\n",
    "# 13 neck\n",
    "annotations = {}\n",
    "# coco annotations info.\n",
    "annotations[\"info\"] = {\n",
    "    \"description\": \"customer dataset format convert to COCO format\",\n",
    "    \"url\": \"http://cocodataset.org\",\n",
    "    \"version\": \"1.0\",\n",
    "    \"year\": 2019,\n",
    "    \"contributor\": \"matt\",\n",
    "    \"date_created\": \"2019/12/19\"\n",
    "}\n",
    "# coco annotations licenses.\n",
    "annotations[\"licenses\"] = [{\n",
    "    \"url\": \"https://www.apache.org/licenses/LICENSE-2.0.html\",\n",
    "    \"id\": 1,\n",
    "    \"name\": \"Apache License 2.0\"\n",
    "}]\n",
    "# coco categories \n",
    "annotations[\"categories\"] = [{\n",
    "    \"supercategory\": \"person\",\n",
    "    \"id\": 1,\"name\": \"person\",\n",
    "    \"keypoints\": [\"right_shoulder\",\"right_elbow\",\"right_wrist\",\n",
    "                  \"left_shoulder\", \"left_elbow\", \"left_wrist\", \n",
    "                  \"right_hip\", \"right_knee\", \"right_ankle\",\n",
    "                  \"left_hip\", \"left_knee\", \"left_ankle\",\n",
    "                 ],\n",
    "    \"skeleton\": [[1,2],[2,3],[4,5],[5,6],[7,8],[8,9],[10,11],[11,12],\n",
    "                 [1,4],[7,10]]\n",
    "}]\n",
    "annotations[\"images\"] = []\n",
    "annotations['annotations'] = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_file = '/media/hsw/E/datasets/ai-challenger/ai_test/ai_val/keypoint_validation_annotations_20170911.json'\n",
    "f = open(json_file, encoding='utf-8')\n",
    "annos = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000\n",
      "10000\n",
      "10000\n",
      "20000\n",
      "20000\n",
      "20000\n",
      "30000\n",
      "30000\n",
      "30000\n",
      "30000\n",
      "30000\n"
     ]
    }
   ],
   "source": [
    "img_path = '/media/hsw/E/datasets/ai-challenger/ai_test/ai_val/val'\n",
    "img_id = 0\n",
    "ann_id = 0\n",
    "coco_pred_annos = []\n",
    "for ann in annos:\n",
    "    img_id += 1\n",
    "#     print (ann)\n",
    "    img = cv2.imread(img_path + \"/\" + ann['image_id'] + \".jpg\")\n",
    "    annotations[\"images\"].append({\n",
    "        \"license\": 1,\n",
    "        \"file_name\": ann['image_id'] + \".jpg\",\n",
    "        \"coco_url\": ann['url'] if 'url' in ann.keys() else '',\n",
    "        \"height\": img.shape[0],\n",
    "        \"width\": img.shape[1],\n",
    "        \"data_captured\": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),\n",
    "        \"flickr_url\": ann['url'] if 'url' in ann.keys() else '',\n",
    "        \"id\": img_id\n",
    "    })\n",
    "    for key, value in ann['human_annotations'].items():\n",
    "        ann_id += 1\n",
    "        x1y1x2y2 = value\n",
    "        kps = ann['keypoint_annotations'][key]\n",
    "        nkps = 0\n",
    "        for i in range(12):\n",
    "            if kps[3*i+2] == 1:\n",
    "                kps[3*i+2] = 2\n",
    "                nkps += 1\n",
    "            elif kps[3*i+2] == 2:\n",
    "                kps[3*i+2] = 1\n",
    "                nkps += 1\n",
    "            elif kps[3*i+2] == 3:\n",
    "                kps[3*i+2] = 0\n",
    "            else:\n",
    "                print ('error in kps')\n",
    "        kps = kps[:12*3]\n",
    "        annotations[\"annotations\"].append({\n",
    "            \"id\": ann_id,\n",
    "            \"image_id\": img_id,\n",
    "            \"category_id\": 1,\n",
    "            \"segmentation\":[[value[0], value[1], value[2], value[1], value[2], value[3], value[0], value[3]]],\n",
    "            \"area\": (value[2]-value[0]) * (value[3]-value[1]),\n",
    "            \"bbox\": [value[0], value[1], value[2]-value[0], value[3]-value[1]],\n",
    "            \"iscrowd\": 0,\n",
    "            \"keypoints\":kps,\n",
    "            \"num_keypoints\":nkps\n",
    "                            \n",
    "        })\n",
    "        coco_pred_annos.append({\n",
    "            \"image_id\": img_id,\n",
    "            \"category_id\": 1,\n",
    "            \"keypoints\": kps,\n",
    "            \"score\": 0.9\n",
    "        })\n",
    "        if img_id % 10000 == 0:\n",
    "            print (img_id)\n",
    "#     break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/media/hsw/E/datasets/multipose_with_only_12_body_joints/aival.json', 'w') as fw:\n",
    "    json.dump(annotations, fw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/media/hsw/E/datasets/cocoapi/annotations/ai_val_pred.json', 'w') as fw:\n",
    "    json.dump(coco_pred_annos, fw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
